import numpy as np

'''该模块负责三角追赶法的计算'''


# 定义温度场的Aw,Ap,Ae及D
# Ap*x[i] = Aw*x[i-1] + Ae*x[i+1] + D （化简式子）
# a[i]*x[i-1] + b[i]*x[i] + c[i]*x[i+1] = D （计算式子）
# b[i]*x[i] = -a[i]*x[i-1] + -c[i]*x[i+1] + D

# 界面处扩散系数求解（调和平均数，算术平均数辅助）
def coefficient_average(coe_P, coe_P_other, x_minus, x_plus):
    average = np.empty_like(coe_P, dtype=float)  # 初始化系数矩阵
    x_sum = x_minus + x_plus  # 求和
    for i in range(len(coe_P)):
        if coe_P[i] * coe_P_other[i] == 0:   # 当存在系数零时，采用算术平均数；否则采用调和平均数
            average[i] = (coe_P[i] * x_plus[i] + coe_P_other[i] * x_minus[i]) / x_sum[i]
        else:
            average[i] = x_sum[i] / (x_minus[i] / coe_P[i] + x_plus[i] / coe_P_other[i])
    return average  # 返回平均数


# 左侧系数AW(-a)求解：（计算场的传导系数）
def AW_coefficient(b_P, x_step):
    b_W = np.concatenate(([1], b_P[:-1]))  # bW运算，左侧添加数字1，不影响后续计算，防止除零错误
    xw_minus = 0.5 * x_step  # xw-计算
    xw_plus = np.concatenate(([0], xw_minus[:-1]))  # xw+计算，左侧添加数字0

    bw = coefficient_average(b_P, b_W,
                             xw_minus, xw_plus)  # 两节点内部扩散系数bw的平均数处理

    A_W = bw / (xw_minus + xw_plus)  # AW计算
    A_W[0] = 0  # AW左侧化为零
    return A_W


# 右侧系数AE(-c)求解：（计算场的传导系数）
def AE_coefficient(b_P, x_step):
    b_E = np.concatenate((b_P[1:], [1]))  # bE运算，右侧添加数字1，不影响后续计算，防止除零错误
    xe_minus = 0.5 * x_step  # xe-计算
    xe_plus = np.concatenate((xe_minus[1:], [0]))  # xe+计算，右侧添加数字0

    be = coefficient_average(b_P, b_E,
                             xe_minus, xe_plus)  # 两节点内部扩散系数be的平均数处理

    A_E = be / (xe_minus + xe_plus)  # AE计算
    A_E[-1] = 0  # AE右侧化为零
    return A_E


# Ap0系数求解
def Ap0_coefficient(ap, x_step, time_step):
    coefficient = ap * x_step / time_step
    return coefficient


# 中间系数AP(b)求解（计算场的时间系数，AW，AE）(AP = Ap0 + AW + AE)
def AP_coefficient(Ap0, AW, AE):
    coefficient = Ap0 + AW + AE  # 计算系数AP（b）
    return coefficient


# 右侧确定项系数Ap0*x0求解
def certain_coefficient(Ap0, x0):
    coefficient = Ap0 * x0
    return coefficient


# 耦合项系数求解
def couple_coefficient(cp, x_step, x_2):
    length = len(x_step)  # 获取矩阵长度
    coefficient = np.zeros(length, dtype=float)  # 初始化耦合项矩阵
    # 计算(cw/dx)，(ce/dx)与-(ce/dx+cw/dx)
    left = AW_coefficient(cp, x_step)
    right = AE_coefficient(cp, x_step)
    center = - (left + right)
    # 计算耦合项系数
    coefficient[0] = center[0] * x_2[0] + right[0] * x_2[1]  # 首尾单独计算，中间部分循环
    coefficient[-1] = left[-1] * x_2[-2] + center[-1] * x_2[-1]
    for i in range(1, length - 1):
        coefficient[i] = left[i] * x_2[i - 1] + center[i] * x_2[i] + right[i] * x_2[i + 1]
    return coefficient


# 源项系数求解（默认源项为零）
def source_coefficient(sp):
    coefficient = sp
    return coefficient


# 总系数D的求解（包括确定项，耦合项，源项，后续扩展）
def final_coefficient(cp, Ap0, x0, x_2, x_step, sp):
    coefficient = certain_coefficient(Ap0, x0)
    coefficient += couple_coefficient(cp, x_step, x_2)
    coefficient += source_coefficient(sp)
    return coefficient


# TDMA三角追赶法
def tdma(AW, AP, AE, D):
    # 三角矩阵赋值并转化为浮点型
    a = np.array(-AW).astype(float)
    b = np.array(AP).astype(float)
    c = np.array(-AE).astype(float)
    d = np.array(D).astype(float)
    length = len(d)  # 得到矩阵长度
    # 三角矩阵首尾置零
    a[0] = 0
    c[-1] = 0
    # 开始计算
    for i in range(1, length):  # 消元
        a[i] = a[i] / b[i - 1]
        b[i] = b[i] - a[i] * c[i - 1]
        d[i] = d[i] - a[i] * d[i - 1]

    x = np.zeros(length).astype(float)  # 初始化x[]
    x[-1] = d[-1] / b[-1]  # 回代计算
    for i in range(length - 2, -1, -1):  # 遍历回代
        x[i] = (d[i] - c[i] * x[i + 1]) / b[i]
    return x


def compute(ap, bp, cp, x_step, time_step, x0, x_2, sp):
    AW = AW_coefficient(bp, x_step)
    AE = AE_coefficient(bp, x_step)
    Ap0 = Ap0_coefficient(ap, x_step, time_step)
    AP = AP_coefficient(Ap0, AW, AE)
    D = final_coefficient(cp, Ap0, x0, x_2, x_step, sp)
    outputs = tdma(AW, AP, AE, D)
    return outputs


def main():
    bp = np.array([1, 0, 3, 1])
    ap = np.array([1, 1, 1, 1])

    x_step = np.array([1, 2, 2, 2])

    dt = 1

    AW = AW_coefficient(bp, x_step)
    AE = AE_coefficient(bp, x_step)
    Ap0 = Ap0_coefficient(ap, x_step, dt)
    AP = AP_coefficient(Ap0, AW, AE)

    print(AW)
    print(AE)
    print(AP)


if __name__ == '__main__':
    main()
